from datetime import datetime

from sqlalchemy import delete, select, update, desc, or_, func
from sqlalchemy.ext.asyncio import AsyncSession

from module_stock.entity.do.settlement_do import Settlement
from module_stock.entity.vo.settlement_vo import SettlementModel, SettlementPageQueryModel
from utils.page_util import PageUtil


class SettlementDao:
    """
    交易交割记录模块数据库操作层
    """

    buy_operators = ['证券买入', '融资买入', '新股入账', '质押回购拆出', '担保品划入', '红股入账']
    sell_operators = ['证券卖出', '拆出质押购回', '股息入账', '股息红利税补', '担保品划出']

    @classmethod
    async def get_settlement_detail_by_id(cls, db: AsyncSession, id: int):
        """
        根据主键ID获取交易交割记录详细信息

        :param db: orm对象
        :param id: 主键ID
        :return: 交易交割记录信息对象
        """
        settlement_info = (
            (
                await db.execute(
                    select(Settlement)
                    .where(
                        Settlement.id == id
                    )
                )
            )
            .scalars()
            .first()
        )

        return settlement_info

    @classmethod
    async def get_settlement_list(cls, db: AsyncSession, query_object: SettlementPageQueryModel, is_page: bool = False):
        """
        根据查询参数获取交易交割记录列表信息

        :param db: orm对象
        :param query_object: 查询参数对象
        :param is_page: 是否开启分页
        :return: 交易交割记录列表信息对象
        """
        if query_object.account_nos:
            query_object.account_no_list = [item.strip() for item in query_object.account_nos.split(',') if
                                            item.strip()]

        query = (
            select(Settlement)
            .where(
                Settlement.trade_date >= datetime.strptime(query_object.begin_trade_date,
                                                           '%Y-%m-%d') if query_object.begin_trade_date else True,
                Settlement.trade_date <= datetime.strptime(query_object.end_trade_date,
                                                           '%Y-%m-%d') if query_object.end_trade_date else True,
                Settlement.create_time == query_object.create_time if query_object.create_time else True,
                Settlement.stock_code == query_object.stock_code if query_object.stock_code else True,
                Settlement.stock_name.like(f'%{query_object.stock_name}%') if query_object.stock_name else True,
                # 查询单个账户
                Settlement.account_no == query_object.account_no if query_object.account_no else True,
                # 同时查询多个账户
                Settlement.account_no.in_(query_object.account_no_list) if query_object.account_no_list else True,
                # 查询单个操作
                Settlement.operation.like(f'%{query_object.operation}%') if query_object.operation else True,
                # 同时查询多个操作
                (or_(*[Settlement.operation == op for op in
                       query_object.operation_list])) if query_object.operation_list else True,
            )
            .order_by(desc(Settlement.trade_date))
            .distinct()
        )
        settlement_list = await PageUtil.paginate(db, query, query_object.page_num, query_object.page_size, is_page)

        return settlement_list

    @classmethod
    async def add_settlement_dao(cls, db: AsyncSession, settlement: SettlementModel):
        """
        新增交易交割记录数据库操作

        :param db: orm对象
        :param settlement: 交易交割记录对象
        :return:
        """
        db_settlement = Settlement(**settlement.model_dump(exclude={}))
        db.add(db_settlement)
        await db.flush()

        return db_settlement

    @classmethod
    async def edit_settlement_dao(cls, db: AsyncSession, settlement: dict):
        """
        编辑交易交割记录数据库操作

        :param db: orm对象
        :param settlement: 需要更新的交易交割记录字典
        :return:
        """
        await db.execute(update(Settlement), [settlement])

    @classmethod
    async def delete_settlement_dao(cls, db: AsyncSession, settlement: SettlementModel):
        """
        删除交易交割记录数据库操作

        :param db: orm对象
        :param settlement: 交易交割记录对象
        :return:
        """
        await db.execute(delete(Settlement).where(Settlement.id.in_([settlement.id])))

    @classmethod
    async def get_unique_stock(cls, db: AsyncSession, account_nos: list, start_date: str):
        """
        获取从指定成交日期开始，所有证券代码不为空的唯一证券代码列表

        :param db: orm对象
        :param account_nos: 账号列表
        :param start_date: 起始成交日期，格式为 'YYYY-MM-DD'
        :return: 唯一证券代码列表
        """
        operation_list = cls.buy_operators + cls.sell_operators
        query = (
            select(Settlement.stock_code, Settlement.stock_name)
            .where(
                Settlement.trade_date <= datetime.strptime(start_date, '%Y-%m-%d') if start_date else True,
                Settlement.stock_code.isnot(None),
                Settlement.stock_code != '',
                Settlement.account_no.in_(account_nos) if account_nos else True,
                (or_(*[Settlement.operation == op for op in operation_list])) if operation_list else True,
            )
            .order_by(Settlement.stock_code)
            .distinct(Settlement.stock_code)
        )
        result = await db.execute(query)
        stock_info_dict = {}
        for row in result.fetchall():
            stock_info_dict[row[0]] = {'stock_code': row[0], 'stock_name': row[1] if row[1] else ''}
        stock_info_list = list(stock_info_dict.values())
        return stock_info_list

    @classmethod
    async def get_buy_and_sell_settlement_list_dao(cls, query_db: AsyncSession, stock_code: str, account_nos: list,
                                                   stat_date: str = None):
        """
        获取指定证券代码的买入和卖出记录

        :param query_db: orm对象
        :param stock_code: 证券代码
        :param account_nos: 账号列表
        :param stat_date: 统计时间
        :return: 买入记录和卖出记录的元组 (buy_records, sell_records)
        """
        # 计算买入总额、买入总量
        buy_query = SettlementPageQueryModel(
            stockCode=stock_code,
            accountNoList=account_nos,
            operationList=cls.buy_operators,
            endTradeDate=stat_date
        )
        buy_records = await SettlementDao.get_settlement_list(query_db, buy_query, False)

        sell_query = SettlementPageQueryModel(
            stockCode=stock_code,
            accountNoList=account_nos,
            operationList=cls.sell_operators,
            endTradeDate=stat_date
        )
        sell_records = await SettlementDao.get_settlement_list(query_db, sell_query, False)

        return buy_records, sell_records

    @classmethod
    async def get_cumulative_trade_amount(cls, db: AsyncSession, account_list: list, stat_date: str):
        """
        获取指定账户在指定日期之前的累计发生金额总和

        :param db: orm对象
        :param account_list: 账号列表
        :param stat_date: 统计截止日期
        :return: 累计发生金额总和
        """
        result = await db.execute(
            select(func.sum(Settlement.occurred_amount))
            .where(
                Settlement.account_no.in_(account_list) if account_list else True,
                Settlement.trade_date <= datetime.strptime(stat_date, '%Y-%m-%d') if stat_date else True,
            )
        )
        amount = result.scalar()
        return amount if amount else 0
